import glob
import os.path as osp
from utils.utils import unpickle
import numpy as np
from scipy.io import loadmat

from .base import BaseDataset


class ssmduke(BaseDataset):
    def __init__(self, root, transforms, split):
        self.name = "ssm_with_duke"
        self.img_prefix = osp.join(root, "Image")
        super(ssmduke, self).__init__(root, transforms, split)

    def _load_split_img_names(self):
        """
        Load the image names for the specific split.
        """
        assert self.split in ("train", "gallery")
        # all images
        all_imgs = glob.glob(osp.join(self.img_prefix,'*.jpg'))
        training_imgs = all_imgs
        return training_imgs

    def _load_annotations(self):

        ssmduke_anno=unpickle('ssmduke/ssm_duke_anno.pkl')
        print('cocops pids: 702, images: ',len(ssmduke_anno))
        return ssmduke_anno
